import random
import argparse

def str2bool(v):
    # susendberg's function
    return v.lower() in ("yes", "true", "t", "1")

def random_select_file_rows(ori_file, w_file, contain, maxrownum=10000, header=True, shuffle_output=False):
    """
    Select rows randomly from the file.
    """
    nrows = 0
    pos_lines = []
    neg_lines = []
    with open(ori_file) as rf:
        for line in rf:
            if contain:
                words = line.strip().split()
                label = int(words[11])
                if label == 0:
                    neg_lines.append(nrows)
                elif label == 1:
                    pos_lines.append(nrows)
            nrows += 1
    if header:
        nrows -= 1
    print('There are {} lines (excluding header) in the file {}'.format(nrows, ori_file))

    actual_nline = maxrownum
    if nrows <= actual_nline:
        actual_nline = nrows
        print('Returning all lines from the file {}'.format(ori_file))
    if not contain:
        random_lines = random.sample(range(1, nrows + 1), actual_nline)
        random_lines = [0] + sorted(random_lines)
        random_lines[-1] = nrows
    else:
        sample_pos = min(int(actual_nline / 2), len(pos_lines))
        sample_neg = actual_nline - sample_pos
        random_lines = random.sample(neg_lines, sample_neg) + random.sample(pos_lines, sample_pos)
        random_lines = sorted(random_lines)

    wf = open(w_file, 'w')
    lines_to_write = []
    with open(ori_file) as rf:
        if header:
            lines_to_write.append(next(rf))
        for i in range(1, len(random_lines)):
            chosen_line = ''
            for j in range(0, random_lines[i] - random_lines[i - 1]):
                chosen_line = next(rf)
            lines_to_write.append(chosen_line)

    if shuffle_output:
        random.shuffle(lines_to_write)

    wf.writelines(lines_to_write)
    wf.close()
    print('random_select_file_rows finished.')

def main():
    parser = argparse.ArgumentParser(description='Select certain number of lines from a file randomly.')
    parser.add_argument('--ori_filepath', type=str, required=True,
                        help='The path of the file where lines are to be selected.')
    parser.add_argument('--write_filepath', type=str, required=True,
                        help='The file to write the selected lines.')
    parser.add_argument('--num_lines', type=int, required=True,
                        help='Number of lines to select.')
    parser.add_argument('--header', type=str, required=False,
                        default='true',
                        help='Whether the original file has a header (true/false). Default is true.')
    parser.add_argument('--contain', action='store_true', default=False,
                        help="Include negative and positive samples.")
    parser.add_argument('--random', action='store_true', default=False,
                        help="Randomly shuffle the output order.")

    args = parser.parse_args()

    orifile = args.ori_filepath
    wfile = args.write_filepath
    srownum = args.num_lines
    header = str2bool(args.header)
    random_select_file_rows(orifile, wfile, args.contain, srownum, header, args.random)

if __name__ == '__main__':
    main()
